import requests
import json
from time import sleep
from typing import List, Dict, Tuple


class TrafficStatusCollector:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://restapi.amap.com/v3/traffic/status/rectangle"

    def split_rectangle(self, start_point: Tuple[float, float], end_point: Tuple[float, float]) -> List[
        Tuple[Tuple[float, float], Tuple[float, float]]]:
        """将矩形区域分为4个子区域"""
        x1, y1 = start_point
        x2, y2 = end_point

        # 计算中间点
        mid_x = (x1 + x2) / 2
        mid_y = (y1 + y2) / 2

        # 返回四个子矩形的坐标
        return [
            ((x1, y1), (mid_x, mid_y)),  # 左下
            ((mid_x, y1), (x2, mid_y)),  # 右下
            ((x1, mid_y), (mid_x, y2)),  # 左上
            ((mid_x, mid_y), (x2, y2))  # 右上
        ]

    def get_traffic_status(self, rectangle: Tuple[Tuple[float, float], Tuple[float, float]], level: int = 5) -> Dict:
        """获取指定矩形区域的交通态势数据"""
        start_point, end_point = rectangle
        rectangle_str = f"{start_point[0]},{start_point[1]};{end_point[0]},{end_point[1]}"

        params = {
            'key': self.api_key,
            'level': level,
            'rectangle': rectangle_str,
            'extensions': 'all',
            'output': 'JSON'
        }

        try:
            response = requests.get(self.base_url, params=params)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"请求失败: {e}")
            return None

    def collect_and_save(self, start_point: Tuple[float, float], end_point: Tuple[float, float],
                         output_file: str, level: int = 4):
        """收集所有子区域的数据并保存为JSON文件"""
        # 分割区域
        sub_rectangles = self.split_rectangle(start_point, end_point)

        # 收集所有子区域的数据
        all_data = []
        for i, rectangle in enumerate(sub_rectangles):
            print(f"正在获取第 {i + 1}/4 个子区域的数据...")
            data = self.get_traffic_status(rectangle, level)
            if data:
                all_data.append(data)
            sleep(0.5)  # 添加延时以避免请求过于频繁

        # 将数据保存为JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump({
                'total_areas': len(all_data),
                'areas_data': all_data
            }, f, ensure_ascii=False, indent=2)

        print(f"数据已保存至 {output_file}")


def main():
    # 使用示例
    api_key = "*****"  # 替换为您的API密钥

    # 定义查询区域的经纬度范围
    start_point = (106.4847586, 29.5319042)  # 左下角坐标
    end_point = (106.5991668, 29.5792822)  # 右上角坐标

    collector = TrafficStatusCollector(api_key)
    collector.collect_and_save(
        start_point=start_point,
        end_point=end_point,
        output_file='traffic_status5.json',
        level=5
    )


if __name__ == "__main__":
    main()